import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn,FasterRCNN
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models import densenet121
from torchvision.models.detection.rpn import AnchorGenerator
from backbone import EfficientDetBackbone
import torch
# 必须使用该方法下载模型，然后加载
from flyai.utils import remote_helper
from efficientnet_pytorch import EfficientNet

def get_model(num_classes = 2):
    model = fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

def get_model_denmodel(num_classes = 2):
    #backbone = densenet121(pretrained=True).features
    denmodel = densenet121(pretrained=False)
    path = remote_helper.get_remote_data('https://www.flyai.com/m/densenet121-a639ec97.pth')
    from torchvision.models.densenet import _load_state_dict 
    _load_state_dict(denmodel,path,progress=True)
    backbone = denmodel.features
    backbone.out_channels = 1024
    anchor_generator = AnchorGenerator(sizes=(16,32,64,128,256),aspect_ratios=(0.5,1,1.5,2))
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],output_size=7,sampling_ratio=2)
    model = FasterRCNN(backbone,num_classes=num_classes,rpn_anchor_generator=anchor_generator,box_roi_pool=roi_pooler)
    return model

def get_model_efficientdet(num_classes=2,is_training=True):
    anchor_ratios = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)]
    anchor_scales = [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]
    #model = EfficientDetBackbone(compound_coef=0, num_classes=90,
    #                         ratios=anchor_ratios, scales=anchor_scales)
    model = EfficientDetBackbone(compound_coef=0, num_classes=2,
                             ratios=anchor_ratios, scales=anchor_scales)
    #if is_training:
    #    weight_path = "/home/gg/桌面/TBDetection_FlyAI-eff/weights/efficientdet-d0.pth"
    #    model.load_state_dict(torch.load(weight_path))
    #model.num_classes = 2
    #print(model)
    return model


if __name__ == "__main__":
    print(get_model_efficientdet())